package com.ehsunbehravesh.service.security.servlets;

import com.ehsunbehravesh.service.security.ejb.encryption.EncryptionLocalBusiness;
import com.ehsunbehravesh.utils.Base64Coder;
import java.io.ByteArrayInputStream;
import java.io.IOException;
import java.io.InputStream;
import java.util.logging.Level;
import java.util.logging.Logger;
import javax.ejb.EJB;
import javax.servlet.ServletException;
import javax.servlet.ServletOutputStream;
import javax.servlet.annotation.MultipartConfig;
import javax.servlet.annotation.WebServlet;
import javax.servlet.http.HttpServlet;
import javax.servlet.http.HttpServletRequest;
import javax.servlet.http.HttpServletResponse;
import javax.servlet.http.Part;

/**
 *
 * @author ehsun7b
 */
@WebServlet(name = "AESFileEncryptionServlet", urlPatterns = {"/AESFileEncryptor"})
@MultipartConfig
public class AESFileEncryptionServlet extends HttpServlet {

  private static final Logger logger = Logger.getLogger(AESFileEncryptionServlet.class.getName());
  private static final int BUFSIZE = 4096;
  @EJB
  private EncryptionLocalBusiness encryptionBusiness;

  public AESFileEncryptionServlet() {
    super();
  }

  @Override
  protected void doPost(HttpServletRequest request, HttpServletResponse response)
          throws ServletException, IOException {

    ServletOutputStream outStream = response.getOutputStream();
    String mimeType = null;
    String errorMessage = null;
    String filename = null;
    byte[] encrypted = null;
    boolean error = false;

    try {

      String keyStr = request.getParameter("key");
      byte[] key = Base64Coder.decode(keyStr);

      for (Part part : request.getParts()) {
        if (!part.getName().equals("file")) {
          continue;
        }
        logger.log(Level.INFO, "--------------");
        logger.log(Level.INFO, "part: {0}", part.getName());
        InputStream is = request.getPart(part.getName()).getInputStream();
        int i = is.available();
        byte[] b = new byte[i];
        is.read(b);
        logger.log(Level.INFO, "Length : {0}", b.length);
        filename = getFileName(part);
        logger.log(Level.INFO, "File name : {0}", filename);

        mimeType = part.getContentType();
        encrypted = encryptionBusiness.encryptAES(key, b);

        is.close();
      }

    } catch (Exception ex) {
      error = true;
      errorMessage = "Fatal error! Please double check your key.";
      logger.log(Level.SEVERE, ex.getMessage());
    }

    if (error) {
      response.setContentType("text/plain");
      byte[] bytes = errorMessage.getBytes("UTF-8");
      response.setContentLength(bytes.length);
      response.setHeader("Content-Disposition", "attachment; filename=\"error\"");

      outStream.write(bytes);
      outStream.close();
    } else {

      response.setContentType(mimeType);
      response.setContentLength(encrypted.length);
      response.setHeader("Content-Disposition", "attachment; filename=\"" + filename + "_encrypted\"");

      ByteArrayInputStream ins = new ByteArrayInputStream(encrypted);
      byte[] byteBuffer = new byte[BUFSIZE];
      int read = 0;
      while ((read = ins.read(byteBuffer)) != -1) {
        outStream.write(byteBuffer, 0, read);
      }

      ins.close();
      outStream.close();
    }
  }

  private String getFileName(Part part) {
    String partHeader = part.getHeader("content-disposition");
    logger.log(Level.INFO, "Part Header = {0}", partHeader);
    for (String cd : part.getHeader("content-disposition").split(";")) {
      if (cd.trim().startsWith("filename")) {
        return cd.substring(cd.indexOf('=') + 1).trim()
                .replace("\"", "");
      }
    }
    return null;

  }
}
